"""
可视化工具模块

提供算法性能和结果的可视化功能
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from typing import Dict, List, Optional, Tuple
import seaborn as sns


class Visualizer:
    """可视化工具类"""
    
    def __init__(self, figsize: Tuple[int, int] = (12, 8)):
        """
        初始化可视化工具
        
        参数:
            figsize: 图形大小
        """
        self.figsize = figsize
        plt.style.use('seaborn-v0_8-darkgrid')
        
    def plot_convergence(self, results: Dict[str, List[float]], 
                        title: str = "收敛性分析",
                        xlabel: str = "迭代次数",
                        ylabel: str = "值",
                        log_scale: bool = False):
        """
        绘制收敛曲线
        
        参数:
            results: 结果字典，键为算法名称，值为迭代值列表
            title: 图标题
            xlabel: x轴标签
            ylabel: y轴标签
            log_scale: 是否使用对数刻度
        """
        plt.figure(figsize=self.figsize)
        
        for algorithm, values in results.items():
            iterations = range(len(values))
            plt.plot(iterations, values, label=algorithm, linewidth=2)
            
        plt.xlabel(xlabel, fontsize=12)
        plt.ylabel(ylabel, fontsize=12)
        plt.title(title, fontsize=14)
        plt.legend(fontsize=10)
        plt.grid(True, alpha=0.3)
        
        if log_scale:
            plt.yscale('log')
            
        plt.tight_layout()
        plt.show()
        
    def plot_dual_convergence(self, primal_results: Dict[str, List[float]],
                             dual_results: Dict[str, List[float]],
                             title: str = "原始-对偶收敛性"):
        """绘制原始和对偶变量的收敛曲线"""
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(self.figsize[0], self.figsize[1]//2))
        
        # 原始变量收敛
        for algorithm, values in primal_results.items():
            iterations = range(len(values))
            ax1.plot(iterations, values, label=algorithm, linewidth=2)
            
        ax1.set_xlabel("迭代次数", fontsize=12)
        ax1.set_ylabel("原始目标值", fontsize=12)
        ax1.set_title("原始变量收敛", fontsize=14)
        ax1.legend(fontsize=10)
        ax1.grid(True, alpha=0.3)
        
        # 对偶变量收敛
        for algorithm, values in dual_results.items():
            iterations = range(len(values))
            ax2.plot(iterations, values, label=algorithm, linewidth=2)
            
        ax2.set_xlabel("迭代次数", fontsize=12)
        ax2.set_ylabel("对偶间隙", fontsize=12)
        ax2.set_title("对偶变量收敛", fontsize=14)
        ax2.legend(fontsize=10)
        ax2.grid(True, alpha=0.3)
        ax2.set_yscale('log')
        
        plt.suptitle(title, fontsize=16)
        plt.tight_layout()
        plt.show()
        
    def plot_constraint_violation(self, inequality_violations: Dict[str, List[float]],
                                 equality_violations: Dict[str, List[float]],
                                 title: str = "约束违反度"):
        """绘制约束违反度"""
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(self.figsize[0], self.figsize[1]))
        
        # 不等式约束违反
        for algorithm, values in inequality_violations.items():
            iterations = range(len(values))
            ax1.semilogy(iterations, values, label=algorithm, linewidth=2)
            
        ax1.set_xlabel("迭代次数", fontsize=12)
        ax1.set_ylabel("不等式约束违反度", fontsize=12)
        ax1.set_title("不等式约束违反", fontsize=14)
        ax1.legend(fontsize=10)
        ax1.grid(True, alpha=0.3)
        
        # 等式约束违反
        for algorithm, values in equality_violations.items():
            iterations = range(len(values))
            ax2.semilogy(iterations, values, label=algorithm, linewidth=2)
            
        ax2.set_xlabel("迭代次数", fontsize=12)
        ax2.set_ylabel("等式约束违反度", fontsize=12)
        ax2.set_title("等式约束违反", fontsize=14)
        ax2.legend(fontsize=10)
        ax2.grid(True, alpha=0.3)
        
        plt.suptitle(title, fontsize=16)
        plt.tight_layout()
        plt.show()
        
    def plot_consensus_error(self, consensus_errors: Dict[str, List[float]],
                            title: str = "共识误差"):
        """绘制共识误差"""
        plt.figure(figsize=self.figsize)
        
        for algorithm, values in consensus_errors.items():
            iterations = range(len(values))
            plt.semilogy(iterations, values, label=algorithm, linewidth=2)
            
        plt.xlabel("迭代次数", fontsize=12)
        plt.ylabel("共识误差 ||y_i - y_avg||", fontsize=12)
        plt.title(title, fontsize=14)
        plt.legend(fontsize=10)
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()
        
    def plot_network_state(self, network, node_values: Dict[int, float],
                          title: str = "网络状态"):
        """可视化网络节点状态"""
        import networkx as nx
        
        plt.figure(figsize=(10, 8))
        
        # 设置节点颜色
        values = [node_values.get(i, 0) for i in range(network.num_nodes)]
        
        # 使用spring布局
        pos = nx.spring_layout(network.graph)
        
        # 绘制网络
        nx.draw_networkx_nodes(network.graph, pos, 
                              node_color=values,
                              cmap='viridis',
                              node_size=500,
                              vmin=min(values),
                              vmax=max(values))
        nx.draw_networkx_labels(network.graph, pos, font_size=12)
        nx.draw_networkx_edges(network.graph, pos, alpha=0.5)
        
        # 添加颜色条
        sm = plt.cm.ScalarMappable(cmap='viridis', 
                                   norm=plt.Normalize(vmin=min(values), 
                                                     vmax=max(values)))
        sm.set_array([])
        plt.colorbar(sm, label="节点值")
        
        plt.title(title, fontsize=16)
        plt.axis('off')
        plt.tight_layout()
        plt.show()
        
    def animate_algorithm(self, network, node_states: List[Dict[int, float]],
                         title: str = "算法动画", interval: int = 100):
        """创建算法执行的动画"""
        import networkx as nx
        
        fig, ax = plt.subplots(figsize=(10, 8))
        
        # 设置布局
        pos = nx.spring_layout(network.graph)
        
        # 初始化
        nodes = nx.draw_networkx_nodes(network.graph, pos, 
                                      node_color=[0]*network.num_nodes,
                                      cmap='viridis',
                                      node_size=500,
                                      vmin=0, vmax=1)
        nx.draw_networkx_labels(network.graph, pos, font_size=12)
        nx.draw_networkx_edges(network.graph, pos, alpha=0.5)
        
        ax.set_title(title, fontsize=16)
        ax.axis('off')
        
        def update(frame):
            if frame < len(node_states):
                values = [node_states[frame].get(i, 0) for i in range(network.num_nodes)]
                nodes.set_array(np.array(values))
                ax.set_title(f"{title} - 迭代 {frame}", fontsize=16)
            return nodes,
            
        anim = FuncAnimation(fig, update, frames=len(node_states),
                           interval=interval, blit=True)
        
        plt.tight_layout()
        return anim
        
    def plot_performance_comparison(self, metrics: Dict[str, Dict[str, float]],
                                   title: str = "算法性能比较"):
        """绘制算法性能比较的柱状图"""
        algorithms = list(metrics.keys())
        metric_names = list(next(iter(metrics.values())).keys())
        
        n_groups = len(metric_names)
        n_bars = len(algorithms)
        
        fig, ax = plt.subplots(figsize=self.figsize)
        
        index = np.arange(n_groups)
        bar_width = 0.8 / n_bars
        
        for i, algorithm in enumerate(algorithms):
            values = [metrics[algorithm][metric] for metric in metric_names]
            ax.bar(index + i * bar_width, values, bar_width,
                  label=algorithm, alpha=0.8)
            
        ax.set_xlabel('性能指标', fontsize=12)
        ax.set_ylabel('值', fontsize=12)
        ax.set_title(title, fontsize=14)
        ax.set_xticks(index + bar_width * (n_bars - 1) / 2)
        ax.set_xticklabels(metric_names)
        ax.legend()
        ax.grid(True, alpha=0.3, axis='y')
        
        plt.tight_layout()
        plt.show()
        
    def plot_parameter_sensitivity(self, param_values: List[float],
                                  results: Dict[str, List[float]],
                                  param_name: str = "参数",
                                  metric_name: str = "性能指标",
                                  title: str = "参数敏感性分析"):
        """绘制参数敏感性分析图"""
        plt.figure(figsize=self.figsize)
        
        for algorithm, values in results.items():
            plt.plot(param_values, values, label=algorithm, 
                    marker='o', linewidth=2, markersize=8)
            
        plt.xlabel(param_name, fontsize=12)
        plt.ylabel(metric_name, fontsize=12)
        plt.title(title, fontsize=14)
        plt.legend(fontsize=10)
        plt.grid(True, alpha=0.3)
        
        # 如果参数值跨度很大，使用对数刻度
        if max(param_values) / min(param_values) > 100:
            plt.xscale('log')
            
        plt.tight_layout()
        plt.show()
        
    def save_all_plots(self, results: Dict, filename_prefix: str = "duca"):
        """保存所有图表"""
        # 这里可以添加保存图表的逻辑
        pass 